#include "armblock.h"

#include "arminst.h"

using namespace Utils;

void ArmBlock::AddInst(Arm::Inst* inst) {
    inst->parent = this;
    instList.push_back(inst);
}

Arm::Terminator* ArmBlock::GetTerminator() const {
    if (instList.empty()) return nullptr;
    auto ret = instList.back()->As<Arm::Terminator>();
    return ret;
}
std::vector<ArmBlock*> ArmBlock::GetNext() const {
    return GetTerminator()->GetNext();
}
std::vector<ArmBlock*> ArmBlock::GetPrev() const {
    std::vector<ArmBlock*> ret;
    ret.reserve(prev.size());
    for (auto [u, _] : prev) ret.push_back(u);
    return ret;
}

void ArmBlock::Connect(ArmBlock* nxt) {
    if (GetTerminator()) return;
    auto jmp = new Arm::Jump(nxt);
    jmp->parent = this;
    instList.emplace_back(jmp);
    nxt->prev[this] = &jmp->nxt;
}

void ArmBlock::Connect(ArithOp cond, ArmBlock* t, ArmBlock* f) {
    if (GetTerminator()) return;
    auto jmp = new Arm::JumpCond(cond, t, f);
    jmp->parent = this;
    instList.emplace_back(jmp);
    t->prev[this] = &jmp->t;
    f->prev[this] = &jmp->f;
}

ArmBlock* ArmBlock::Remove(int reserve) {
    // Remove the block and connect the previous with the reserved block.
    // reserve = 0/1: reserve true/false condition block.
    if (reserve == -1) reserve = GetNext()[0] == this;
    Assert(reserve < GetNext().size(), "");
    auto v = GetNext()[reserve];
    Assert(v->prev.count(this), "reverse edge not found");
    v->prev.erase(this);
    // u -> (this) -> v
    for (auto [u, ref] : prev) {
        Assert(*ref == this, "reverse edge not pointing to the right node");
        v->prev[u] = ref, *ref = v;
    }
    return v;
}

ArmBlock* ArmBlock::RemoveBranch(int reserve) {
    // Change branch terminator into jump terminator.
    // reserve = 0/1: reserve true/false condition block.
    if (reserve == -1) reserve = GetNext()[0] == this;
    auto w = GetNext()[!reserve];
    Assert(w->prev.count(this), "reverse edge not found");
    w->prev.erase(this);
    auto v = GetNext()[reserve];
    instList.pop_back();
    auto jmp = new Arm::Jump(v);
    jmp->parent = this;
    instList.emplace_back(jmp);
    v->prev[this] = &jmp->nxt;
    return v;
}

void ArmBlock::GenAsm(std::ostream& os) const {
    os << name << ":\n";
    for (auto inst : instList) {
        os << "    ";
        inst->GenAsm(os);
        os << "\n";
    }
}
